"""
training .py for segmentation specific for branch
"""
import keras
from keras.callbacks import TensorBoard
from data_utils import generator
from new_model.build_model import model_generator
import os
from new_model.metrics import my_iou_metric


def create_callbacks():
    callback = []
    checkpoint = keras.callbacks.ModelCheckpoint(
        os.path.join(
            '',
            'segmentation_model'
        ),
        verbose=1
    )
    tensorboard = TensorBoard(log_dir='./log/segmentation')
    callback.append(checkpoint)
    callback.append(tensorboard)
    return callback


def main():
    # build model
    detection_model, segmentation_model, unit_model = model_generator()
    # compile model
    metrics = ['acc', my_iou_metric]
    segmentation_model.compile(
        loss=keras.losses.sparse_categorical_crossentropy,
        optimizer='adam',
        metrics=metrics
    )
    # load weight if exist
    # segmentation_dir = 'D:\herschel\\navigation\keras_implementation\VGG_segmentation_model_on_cityscape'
    segmentation_dir = None
    if (segmentation_dir is not None) and os.path.exists(segmentation_dir):
        segmentation_model.load_weights(segmentation_dir, by_name=True, skip_mismatch=True)
    # train on cityscape
    train_generator = generator.segmentation_generator('train', batch_size=8)
    val_generator = generator.segmentation_generator('val', batch_size=8)
    # train_generator = generator.segmentation_generator_test(batch_size=8)
    # val_generator = generator.segmentation_generator_test(batch_size=8)
    print('start training...')
    callbacks = create_callbacks()
    segmentation_model.fit_generator(generator=train_generator,
                                     steps_per_epoch=1500,
                                     epochs=10,
                                     verbose=1,
                                     callbacks=callbacks,
                                     validation_data=val_generator,
                                     validation_steps=1,
                                     workers=3,
                                     )
    segmentation_model.save('segmentation_model_on_cityscape')


if __name__ == '__main__':
    # import os
    # os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  # see issue #152
    # os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    main()